from abc import ABCMeta, abstractmethod
from .farthest_points_sampling import fps
from .contacts import Contact3D
from .suctionpoints import SuctionPoint
from .appoach_grasp import ApproachParallelJawPtGrasp3D
from .grasp import ParallelJawPtGrasp3D
from .suction import Suction3D
import trimesh
import numpy as np
import logging
import random
import time

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)


class GraspSampler:
    """ Base class for various methods to sample a number of grasps on an object.
    Should not be instantiated directly.

    Attributes
    ----------
    gripper : :obj:`RobotGripper`
        the gripper to compute grasps for
    config : :obj:`YamlConfig`
        configuration for the grasp sampler
    """
    __metaclass__ = ABCMeta

    def __init__(self, gripper, config):
        self.gripper = gripper
        self._configure(config)

    def _configure(self, config):
        """ Configures the grasp generator."""
        self.friction_coef = config['sampling_friction_coef']
        self.num_cone_faces = config['num_cone_faces']
        self.num_samples = config['grasp_samples_per_surface_point']
        self.target_num_grasps = config['target_num_grasps']
        if self.target_num_grasps is None:
            self.target_num_grasps = config['min_num_grasps']

        # self.min_contact_dist = config['min_contact_dist']
        self.num_grasp_rots = config['num_grasp_rots']
        if 'max_num_surface_points' in list(config.keys()):
            self.max_num_surface_points_ = config['max_num_surface_points']
        else:
            self.max_num_surface_points_ = 500
        if 'grasp_dist_thresh' in list(config.keys()):
            self.grasp_dist_thresh_ = config['grasp_dist_thresh']
        else:
            self.grasp_dist_thresh_ = 0
            
    @abstractmethod
    def sample_grasps(self, graspable, num_grasps_generate, points_num, collision_manager, vis, **kwargs):
        """
        Create a list of candidate grasps for a given object.
        Must be implemented for all grasp sampler classes.

        Parameters
        ---------
        graspable : :obj:`GraspableObject3D`
            object to sample grasps on
        num_grasps_generate : int
        vis : bool
        """
        grasp = []
        return grasp

    @abstractmethod
    def nms_grasps(grasps, threshold=0.005):
        pruned_grasps = []
        score = [grasp.grasp_score.sum() for grasp in grasps]
        score = np.array(score)
        order = np.argsort(score)
        while order.size > 0:
            index = order[-1]
            cur_grasp = grasps[index]
            pruned_grasps.append(cur_grasp)
            approach = True
            if approach:
                distance = [ApproachParallelJawPtGrasp3D.distance(cur_grasp, grasps[idx]) for idx in order[:-1]]
            else:
                distance = [ParallelJawPtGrasp3D.distance(cur_grasp, grasps[idx]) for idx in order[:-1]]
            distance = np.array(distance)
            left = np.where(distance > threshold)
            order = order[left]
        return pruned_grasps

    @staticmethod
    def create_collision_manager():
        collision_manager = trimesh.collision.CollisionManager()
        return collision_manager

    @staticmethod
    def add_collision_obj(collision_manager, graspable):
        collision_manager.add_object('{}'.format(graspable.key), graspable.mesh)

    @staticmethod
    def remove_collision_obj(collision_manager, graspable):
        collision_manager.remove_object('{}'.format(graspable.key))


    def generate_grasps(self, graspable, target_num_grasps=None, grasp_gen_mult=5, max_iter=1,
                        sample_approach_angles=False, points_num=500, vis=False, **kwargs):
        """Samples a set of grasps for an object.

        Parameters
        ----------
        graspable : :obj:`GraspableObject3D`
            the object to grasp
        target_num_grasps : int
            number of grasps to return, defualts to self.target_num_grasps
        grasp_gen_mult : int
            number of additional grasps to generate
        max_iter : int
            number of attempts to return an exact number of grasps before giving up
        sample_approach_angles : bool
            whether or not to sample approach angles
        vis : bool
            whether show the grasp on picture

        Return
        ------
        :obj:`list` of :obj:`ParallelJawPtGrasp3D`
            list of generated grasps
        """
        # assert self.collision_manager is not None
        collision_manager = GraspSampler.create_collision_manager()

        GraspSampler.add_collision_obj(collision_manager, graspable)
        # get num grasps
        if target_num_grasps is None:
            target_num_grasps = self.target_num_grasps
        num_grasps_remaining = target_num_grasps

        grasps = []
        k = 1
        while num_grasps_remaining > 0 and k <= max_iter:
            # SAMPLING: generate more than we need
            num_grasps_generate = grasp_gen_mult * num_grasps_remaining
            new_grasps, badpoints_list = self.sample_grasps(graspable, num_grasps_generate, points_num, collision_manager, vis, **kwargs)

            # add to the current grasp set
            grasps += new_grasps

            # COVERAGE REJECTION: prune grasps by distance
#            grasps = GraspSampler.nms_grasps(grasps, threshold=0.005)

            # ANGLE EXPANSION sample grasp rotations around the axis
#            candidate_grasps = []
#            if sample_approach_angles:
#                for grasp in pruned_grasps:
#                    # construct a set of rotated grasps
#                    for i in range(self.num_grasp_rots):
#                        rotated_grasp = copy.copy(grasp)
#                        delta_theta = 0  # add by Hongzhuo Liang
#                        print("This function can not use yes, as delta_theta is not set. --Hongzhuo Liang")
#                        rotated_grasp.set_approach_angle(i * delta_theta)
#                        candidate_grasps.append(rotated_grasp)
#            else:
#                candidate_grasps = pruned_grasps
#
            logger.info('%d/%d grasps found after iteration %d.',
                        len(grasps), target_num_grasps, k)

            grasp_gen_mult *= 2
            num_grasps_remaining = target_num_grasps - len(grasps)
            k += 1

        # shuffle computed grasps
        random.shuffle(grasps)
        logger.info('Found %d grasps.', len(grasps))
        GraspSampler.remove_collision_obj(collision_manager, graspable)
        return grasps, badpoints_list


class AntipodalGraspSampler(GraspSampler):
    """ Samples antipodal pairs using rejection sampling.
    The proposal sampling ditribution is to choose a random point on
    the object surface, then sample random directions within the friction cone,
    then form a grasp axis along the direction,
    close the fingers, and keep the grasp if the other contact point is also in the friction cone.
    """

    def sample_from_cone(self, n, tx, ty, num_samples=1):
        """ Samples directoins from within the friction cone using uniform sampling.

        Parameters
        ----------
        n : 3x1 normalized :obj:`numpy.ndarray`
            surface normal
        tx : 3x1 normalized :obj:`numpy.ndarray`
            tangent x vector
        ty : 3x1 normalized :obj:`numpy.ndarray`
            tangent y vector
        num_samples : int
            number of directions to sample

        Returns
        -------
        v_samples : :obj:`list` of 3x1 :obj:`numpy.ndarray`
            sampled directions in the friction cone
       """
        v_samples = []
        for i in range(num_samples):
            theta = 2 * np.pi * np.random.rand()
            r = self.friction_coef * np.random.rand()
            v = n + r * np.cos(theta) * tx + r * np.sin(theta) * ty
            v = -v / np.linalg.norm(v)
            v_samples.append(v)
        return v_samples

    def perturb_point(self, x, scale):
        """ Uniform random perturbations to a point """
        x_samp = x + (scale / 2.0) * (np.random.rand(3) - 0.5)
        return x_samp

    def sample_grasps(self, graspable, num_grasps, points_num, collision_manager, vis=False, **kwargs):
        """Returns a list of candidate grasps for graspable object.

        Parameters
        ----------
        graspable : :obj:`GraspableObject3D`
            the object to grasp
        num_grasps : int
            number of grasps to sample
        vis : bool
            whether or not to visualize progress, for debugging

        Returns
        -------
        :obj:`list` of :obj:`ParallelJawPtGrasp3D`
            the sampled grasps
        """
        # get sampled surface points
        grasps = []
        #print(self.max_num_surface_points_, points)
        points_num = min(self.max_num_surface_points_, points_num)
        points_surface_num = 4 * points_num
        sampled_points, tri_index = graspable.mesh.sample(points_surface_num, return_index=True)

        #sample_point
        #print('points is : ', points)
        #start = time.time()

        sample_points_index = fps(sampled_points, points_num, graspable.mesh, use_geodesic=False)
        #time_cost = time.time() - start
        #exit()
        #sample_points_index = fps(graspable.mesh, points_num)
        assert len(sample_points_index) == points_num
        #surface_points, _ = graspable.mesh.vertices[sample_points_index]
        #np.random.shuffle(sample_points_index)
        #shuffled_surface_points = surface_points[:min(self.max_num_surface_points_, len(surface_points))]
        logger.info('Num surface: %d' % (len(sample_points_index)))

        badpoint_list = []
        goodpoint_lsit = []
        #for k, x_surf in enumerate(surface_points_index):
        for idx in sample_points_index:
            break_flag = False
            start_time = time.clock()

            neighbors = graspable.mesh.faces[tri_index[idx]]
            #print(neighbors)
            #exit()
            #if not graspable.badpoints_mask[idx]:
            #    badpoint_list.append(idx)
            #    #print('found badpoint {}'.format(idx))
            #    continue

            #neighbors = graspable.mesh.vertex_neighbors[idx]
            for neighbor in neighbors:
                if not graspable.badpoints_mask[neighbor]:
                    badpoint_list.append(sampled_points[idx])
                    break_flag = True
                    break
            if break_flag:
                #print('meet bad point neighbor')
                continue

            # perturb grasp for num samples
            success_tmp = False
            for i in range(self.num_samples):
                # perturb contact (TODO: sample in tangent plane to surface)
                #x1 = self.perturb_point(x_surf, graspable.sdf.resolution)
                # compute friction cone faces

                c1 = Contact3D(graspable, index=tri_index[idx], contact_point=sampled_points[idx], normal=None, in_direction=None)
                #print('************attempt {} time*************'.format(i))
                _, tx1, ty1 = c1.tangents()
                cone_succeeded, cone1, n1 = c1.friction_cone(self.num_cone_faces, self.friction_coef)
                if not cone_succeeded:
                    continue
                cone_time = time.clock()

                # sample grasp axes from friction cone
                v_samples = self.sample_from_cone(n1, tx1, ty1, num_samples=3)
                v_samples.append(-np.asarray(n1))
                sample_time = time.clock()
                v_sample_numbers = 0

                for v in v_samples:
                    v_sample_numbers += 1
                    #print('*****sample {} time*****'.format(v_sample_numbers))
                    if vis:
                        pass

                    # start searching for contacts
                    contact_found, grasp_candidates, c2 = ParallelJawPtGrasp3D.grasp_from_contact_and_axis_on_mesh(
                        self.gripper, graspable, c1, v, self.gripper.max_width,
                        friction_coef=self.friction_coef, vis=vis)
                    if not contact_found:
                        continue
                    #print('found grasp candidate')
                    #print(grasp)
                    #flag = False
                    for grasp, c2_tmp in zip(grasp_candidates, c2):
                        if 'random_approach_angle' in kwargs and kwargs['random_approach_angle']:
                            angle_candidates = np.arange(-90, 120, 30)
                            np.random.shuffle(angle_candidates)
                            for grasp_angle in angle_candidates:
                                grasp.approach_angle_ = grasp_angle
                                # get true contacts (previous is subject to variation)
                                success = grasp.close_fingers(graspable, self.gripper, collision_manager, vis=vis)
                                if not success:
                                    continue
                                break
                            #else:
                            #    continue
                        else:
                            success = grasp.close_fingers(graspable, self.gripper, collision_manager, vis=vis)
                            if not success:
                                continue
                        if success:
                            #print('success found grasp')
                            success_tmp = True
                            
                            ### GWS annalysis ###
                            #cone_succeeded, cone1, n1
                            #forces = np.zeros([3, 0])
                            #torques = np.zeros([3, 0])
                            #normals = np.zeros([3, 0])
                            #for contact in [c1, c2_tmp]:
                            #    force_success, contact_forces, contact_outward_normal = contact.friction_cone(self.num_cone_faces, self.friction_coef)
                            #    if not force_success:
                            #        print('Force computation failed')
                            #        exit()
                            #    torque_success, contact_torques = contact.torques(contact_forces)
                            #    if not torque_success:
                            #        print('Torque computation failed')
                            #        exit()
                            #    n = contact.normal_force_magnitude()
                            #    forces = np.c_[forces, n * contact_forces]
                            #    torques = np.c_[torques, n * contact_torques]
                            #    normals = np.c_[normals, n * -contact_outward_normal]
                            #if normals.shape[1] == 0:
                            #    print('No normals')
                            #    exit()
                            #if 'torque_scaling' not in params.keys():
                            #    torque_scaling = 1.0
                            #    if method == 'ferrari_canny_L1':
                            #        mn = np.min(graspable.mesh.vertices, axis=0)
                            #        mx = np.max(graspable.mesh.vertices, axis=0)
                            #        torque_scaling = 1.0 / np.median(mx)
                            #    params.torque_scaling = torque_scaling
                            #Q_func = getattr(PointGraspMetrics3D, params.quality_method)
                            #quality = Q_func(forces, torques, normals, soft_fingers=params.soft_fingers,params=params)
                        
                            #if quality > 0.25:
                            #    print(quality)
                            #    print(grasp.grasp_score)
                            #    success = grasp.close_fingers(graspable, self.gripper, self.collision_manager, vis=True)
                            ### GWS annalysis ###
                            #if grasp.grasp_score > 0.5:
                            #    success = grasp.close_fingers(graspable, self.gripper, self.collision_manager, vis=True)

                            grasps.append(grasp)

            if not success_tmp:
                badpoint_list.append(sampled_points[idx])
        random.shuffle(grasps)
        return grasps, badpoint_list


class SuctionSampler(GraspSampler):
    def sample_grasps(self, graspable, num_grasps, points_num, collision_manager, vis=False, **kwargs):

        # get sampled surface points
        grasps = []
        # print(self.max_num_surface_points_, points)
        area = graspable.mesh.area
        k = int(area / 3e-6)
        points_num = min(self.max_num_surface_points_, points_num, k)
        points_surface_num = 3 * points_num
        sampled_points, tri_index = graspable.mesh.sample(points_surface_num, return_index=True)
        
        sample_points_index = fps(sampled_points, points_num, graspable.mesh, use_geodesic=False)

        assert len(sample_points_index) == points_num
        logger.info('Num surface: %d' % (len(sample_points_index)))

        badpoint_list = []
        for idx in sample_points_index:
            break_flag = False
            # fileter unreasonable suction point with potential badpoints
            neighbors = graspable.mesh.faces[tri_index[idx]]
            for neighbor in neighbors:
                if not graspable.badpoints_mask[neighbor]:
                    badpoint_list.append(sampled_points[idx])
                    break_flag = True
                    break
            if break_flag:
                # print('meet bad point neighbor')
                continue

            # perturb grasp for num samples
            success_tmp = False

            c1 = SuctionPoint(graspable, index=tri_index[idx], suction_point=sampled_points[idx], normal=None, in_direction=None)

            v_samples = [-c1.normal]
            v_sample_numbers = 0

            for v in v_samples:
                v_sample_numbers += 1
                # print('*****sample {} time*****'.format(v_sample_numbers))
                if vis:
                    pass

                # start searching for contacts
                success_tmp, grasp_candidate, R = Suction3D.suction_from_point_and_axis_on_mesh(graspable, c1, v,
                                                                                                vis=vis)
                # if not success then go continue
                if not success_tmp:
                    continue

                success_tmp = grasp_candidate.close_fingers(graspable, self.gripper, collision_manager, R, vis=vis)

                if success_tmp:
                    grasps.append(grasp_candidate)
                    break

            if not success_tmp:
                badpoint_list.append(sampled_points[idx])

        random.shuffle(grasps)
        return grasps, badpoint_list


class ApproachAntipodalGraspSampler(GraspSampler):
    def sample_grasps(self, graspable, num_grasps, points_num, collision_manager, vis=False, **kwargs):

        # get sampled surface points
        grasps = []
        area = graspable.mesh.area
        k = int(area / 5e-6)
        points_num = min(self.max_num_surface_points_, points_num, k)
        points_surface_num = 8 * points_num
        sampled_points, tri_index = graspable.mesh.sample(points_surface_num, return_index=True)

        sample_points_index = fps(sampled_points, points_num, graspable.mesh, use_geodesic=False)

        assert len(sample_points_index) == points_num
        logger.info('Num surface: %d' % (len(sample_points_index)))

        badpoint_list = []
        vis = False
        if vis:
            pc = trimesh.PointCloud(sampled_points[sample_points_index],colors=[0,255,0])
            scene = trimesh.Scene([graspable.mesh,pc])
            scene.show()

        for idx in sample_points_index:
            success_tmp = False
            # print('self number of points', self.num_samples)
            # exit()

            sampled_point = Contact3D(graspable, index=tri_index[idx], contact_point=sampled_points[idx], normal=None,
                                      in_direction=None)
            v_samples = [-sampled_point.normal]
            for v in v_samples:
                # start searching for contacts
                contact_found, grasp_candidates = ApproachParallelJawPtGrasp3D.grasp_from_point_and_axis_on_mesh(
                    self.gripper, graspable, sampled_point, v, self.gripper.max_width,
                    friction_coef=self.friction_coef, collision_manager=collision_manager, vis=vis)

                if not contact_found:
                    continue
                # print('len of grasp candidates:', len(grasp_candidates))
                # k = 0
                # if len(grasp_candidates) > 0:
                else:
                    success_tmp = True
                    # for grasp in grasp_candidates:
                    grasps.extend(grasp_candidates)
                    # grasp_candidates += grasp_candidates

            if not success_tmp:
                badpoint_list.append(sampled_points[idx])
        random.shuffle(grasps)
        return grasps, badpoint_list
